import os
from pinecone import Pinecone
from sklearn.model_selection import train_test_split
from tqdm import tqdm
import logging
from sklearn.datasets import load_digits
from pinecone import Pinecone, ServerlessSpec

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Load MNIST dataset
digits = load_digits()
X, y = digits.data, digits.target

# Split the dataset
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Initialize Pinecone
api_key = "ecb9978d-a69b-458c-aa23-37d03b1407f7"
environment = "us-east-1"
pc = Pinecone(api_key=api_key)

index_name = "quickstart"

# Check if index exists, if not create it
if index_name not in pc.list_indexes().names():
    try:
        pc.create_index(
            name=index_name,
            dimension=X_train.shape[1],
            metric="cosine",
            spec=ServerlessSpec(
                cloud='aws',
                region='us-east-1'  # 使用提供的区域
            )
        )
        logging.info(f"Created new index: {index_name}")
    except Exception as e:
        logging.error(f"Failed to create index: {e}")
        raise

# Connect to the index
index = pc.Index(index_name)

# Upload training data to Pinecone
batch_size = 100
total_vectors = len(X_train)
with tqdm(total=total_vectors, desc="Uploading vectors") as pbar:
    for i in range(0, total_vectors, batch_size):
        batch = X_train[i:i+batch_size]
        ids = [f"vec_{j}" for j in range(i, min(i+batch_size, total_vectors))]
        vectors = [(ids[j], batch[j].tolist(), {"label": int(y_train[i+j])}) for j in range(len(batch))]
        index.upsert(vectors=vectors)
        pbar.update(len(batch))

logging.info(f"Successfully uploaded {total_vectors} vectors to index {index_name}")

# Function to get KNN prediction
def knn_predict(query_vector, k=11):
    results = index.query(vector=query_vector.tolist(), top_k=k, include_metadata=True)
    votes = [match['metadata']['label'] for match in results['matches']]
    return max(set(votes), key=votes.count)

# Test accuracy with k=11
k = 11
correct = 0
total = len(X_test)

with tqdm(total=total, desc=f"Testing accuracy (k={k})") as pbar:
    for i in range(total):
        prediction = knn_predict(X_test[i], k)
        if prediction == y_test[i]:
            correct += 1
        pbar.update(1)

accuracy = correct / total
logging.info(f"Accuracy with k={k}: {accuracy:.4f}")

# Optionally, delete the index
# Uncomment the following line if you want to delete the index after the script runs
# pc.delete_index(index_name)